"""
author: liuxu
date: 2025/9/20
description: 检测mpi环境是否可用，获取pmc相关信息
"""
import os
import tempfile
from typing import Tuple, Optional
from common.logger_ctrl import mylogger
from common.ssh_command import exec_cmd


def check_mpi_environment(
        mpi_command: str = "mpiexec",
        min_processes: int = 2
) -> Optional[Tuple[bool, dict]]:
    """
    检查当前MPI环境是否可用

    Args:
        mpi_command: MPI执行命令 (mpiexec, mpirun等)
        min_processes: 需要的最小进程数

    Returns:
        Tuple[bool, str]: (是否可用, mpi详细信息)
    """
    # 0、定义需要获取的mpi信息
    mpi_version = None
    mpi_so_abspath = None


    # 1. 检查MPI命令是否存在
    res, version_output = exec_cmd(f"mpiexec --version") or exec_cmd("mpirun --version")
    if res is True:
        mpi_version = version_output.splitlines()[0]
    else:
        mylogger.error(f"MPI命令({mpi_command}/mpirun)未找到或不可用")
        return False, {}
    # 根据mpicc路径获取mpi.so路径
    res, output = exec_cmd("which mpicc")
    mpicc_abspath = output.strip()
    mpi_so_abspath = os.path.join(mpicc_abspath.rsplit("bin", maxsplit=1)[0], "lib/libmpi.so")
    # mpi_exist = os.path.exists(mpi_so_abspath)
    mpi_info_d = {"mpi_version": mpi_version, "mpi_so_abspath": mpi_so_abspath}

    # 2. 创建简单的C语言MPI测试程序
    test_code = '''
#include <mpi.h>
#include <stdio.h>

int main(int argc, char** argv) {
    MPI_Init(NULL, NULL);
    
    int world_rank;
    MPI_Comm_rank(MPI_COMM_WORLD, &world_rank);
    int world_size;
    MPI_Comm_size(MPI_COMM_WORLD, &world_size);
    
    // 简单的通信测试
    if (world_rank == 0) {
        int number = 42;
        for (int i = 1; i < world_size; i++) {
            MPI_Send(&number, 1, MPI_INT, i, 0, MPI_COMM_WORLD);
        }
    } else {
        int number;
        MPI_Recv(&number, 1, MPI_INT, 0, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        printf("Process %d received number %d from process 0\\n", world_rank, number);
    }
    
    MPI_Barrier(MPI_COMM_WORLD);
    
    if (world_rank == 0) {
        printf("MPI测试成功: 共%d个进程\\n", world_size);
    }
    
    MPI_Finalize();
    return 0;
}
'''

    # 将测试代码写入临时文件
    with tempfile.NamedTemporaryFile(mode='w', suffix='.c', delete=False) as f:
        f.write(test_code)
        test_file = f.name

    executable = test_file.replace('.c', '')

    try:
        # 3. 编译MPI程序
        compile_cmd = f"mpicc {test_file} -o {executable}"
        compile_result, compile_output = exec_cmd(compile_cmd.split(), timeout=30)

        if not compile_result:
            mylogger.error(f"MPI程序编译失败: {compile_output}")
            return False, mpi_info_d

        # 4. 尝试运行MPI程序
        run_cmd = [mpi_command, "-n", str(min_processes), executable]
        run_result, run_output = exec_cmd(run_cmd, timeout=30)

        if run_result:
            return True, mpi_info_d
        else:
            mylogger.error(f"The MPI demo failed to run. Please check the environment: {run_output}")
            return False, mpi_info_d

    except Exception as e:
        mylogger.error(f"The MPI environment is unavailable: {str(e)}")
        return False, mpi_info_d
    finally:
        # 清理临时文件
        try:
            os.unlink(test_file)
            if os.path.exists(executable):
                os.unlink(executable)
        except Exception:
            pass


if __name__ == '__main__':
    res=check_mpi_environment()
    print(res)